import abc
from typing import Any, List
from langchain.tools import BaseTool
import logging
import pandas as pd
from utils.configFinRAG import sql_examples_path
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

logging.basicConfig(level=logging.INFO)

class FinSQLRAG(BaseTool, abc.ABC):
    name = "查询金融数据库"
    description = "当被问到金融查询相关的问题时，会去金融数据库检索结果"

    def __init__(self):
        super().__init__()
        self.example_question_list: List[str] = []
        self.example_sql_list: List[str] = []
        self.example_info_list: List[str] = []
        self.example_fa_list: List[str] = []
        self.example_token_list: List[Any] = []
        self._load_examples()

    def _load_examples(self):
        """加载SQL示例数据"""
        examples = pd.read_csv(sql_examples_path, delimiter=",", header=0)
        for _, row in examples.iterrows():
            self.example_question_list.append(row['问题'])
            self.example_sql_list.append(row['SQL'])
            self.example_info_list.append(row['资料'])
            self.example_fa_list.append(row['FA'])
            tokens = TOKENIZER(row['问题'])['input_ids']
            self.example_token_list.append(tokens)

    async def _arun(self, *args: Any, **kwargs: Any) -> Any:
        # 不予具体实现
        pass

    def _run(self, para) -> str:
        query = para
        try:
            result_prompt, sql = generate_sql(query, LLM, self.example_question_list, self.example_sql_list,
                                              self.example_token_list)
            logging.info(f"Query SQL is: {sql}")
            with sqlite3.connect('/path/to/your/database.db') as conn:
                cs = conn.cursor()
                success_flag, exc_result = query_db(sql, cs)
                if not success_flag:
                    raise Exception("Database query failed.")
                answer = generate_answer(query, exc_result, LLM, self.example_question_list, self.example_info_list,
                                         self.example_fa_list, self.example_token_list)
                return answer
        except Exception as e:
            logging.error(f"Retriever Error: {e}")
            return 'FinSQLRAG处理异常！'

if __name__ == "__main__":
    tool = FinSQLRAG()
    query = "请帮我计算，在20210715，中信行业分类划分的一级行业为消费者服务行业中，涨跌幅最大股票的股票代码是？涨跌幅是多少？百分数保留两位小数。股票涨跌幅定义为：（收盘价 - 前一日收盘价 / 前一日收盘价）* 100%。"
    result = tool._run(query)
    print(result)